from data.common_data_process import get_positive_data, consolidate_dataset, get_user_id_list
import numpy as np
import time
import torch
from config import cfg
from model.model import Model
from util import record_result


def print_final_result(results):
    # calculate final results
    results = np.array(results, dtype=np.float64)
    means = np.mean(results, axis=0)
    stds = np.std(results, axis=0)
    # print final results
    record_result('\nFinal accuracy mean: {}'.format(means[0]))
    record_result('Final accuracy std: {}'.format(stds[0]))
    record_result('\nFinal recall mean: {}'.format(means[1]))
    record_result('Final recall std: {}'.format(stds[1]))
    record_result('\nFinal frp mean: {}'.format(means[4]))
    record_result('Final frp std: {}'.format(stds[4]))
    record_result('\nFinal precision mean: {}'.format(means[2]))
    record_result('Final precision std: {}'.format(stds[2]))
    record_result('\nFinal F1 mean: {}'.format(means[3]))
    record_result('Final F1 std: {}'.format(stds[3]))
    record_result("============================================")

if __name__ == "__main__":

    torch.cuda.set_device(0)

    start = int(round(time.time() * 1000))
    #模型初始化
    print("loading model:" + cfg.net)
    model = Model()
    print("The loading is complete")
    model_count = 0
    results = []
    user_id_data = get_user_id_list(cfg.user_id_file_path)
    print("  the count of user :" + str(len(user_id_data)))
    for positive_user_id in user_id_data:
        print('Train model for user: {}'.format(positive_user_id))
        model_count = model_count + 1

        positive_data, positive_labels = get_positive_data(positive_user_id, cfg.user_data_file_path, cfg.user_id_file_path)
        with open(cfg.user_data_file_path, 'r') as f:
            negative_data = []
            negative_labels = []
            count = 1
            batch = 0
            for i, line in enumerate(f):
                if i == positive_user_id:
                    continue
                sample = np.array(line.split(',')).astype(np.float)
                negative_data.append(sample)
                negative_labels.append(0)
                count += 1
                if count > positive_data.shape[0] * cfg.copy_size:
                    batch = batch + 1
                    negative_data = np.array(negative_data)
                    negative_labels = np.array(negative_labels)

                    x_train, x_test, y_train, y_test = consolidate_dataset(positive_data, negative_data, positive_labels, negative_labels,
                                                                           cfg.data_row_size, cfg.step , cfg.word_dict_size)
                    print('------------------------------------------------------------')
                    print('start ' + str(batch) + 'th training')
                    print('positive data_num : negative data_num = ' + str(len(positive_labels) * cfg.copy_size) + ' : ' + str(len(negative_labels)))

                    model.train_model(x_train, y_train, cfg.batch_size, cfg.epochs, positive_user_id, cfg.save_model_path)

                    record_result('***********************************************************')
                    record_result('user:' + str(positive_user_id) + ', ' + str(batch) + 'th training result')

                    result, step_out = model.evaluate_model( x_test, y_test,cfg.show_roc)

                    if step_out:
                        break
                    negative_data = []
                    negative_labels = []
                    count = 1
            results.append(result)
            if model_count >= cfg.max_model_count:
                break
    # calculate final results
    print_final_result(results)
    # calculate total time
    end = int(round(time.time() * 1000))
    print('\nTotal time expended: {}'.format(end - start))
